import os
import zipfile
import sqlite3
import argparse
from pathlib import Path
# from tqdm import tqdm

# 定义图片扩展名集合
image_extensions = {'.jpg', '.jpeg', '.png', '.aae', '.heic', '.thm', '.gif', '.bmp', '.tiff', '.webp'}

def get_photo_info(db_path, photo_hash):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("SELECT zip_path, zip_name, file_path, file_name FROM photo_hash WHERE file_hash=?", (photo_hash,))
    return c.fetchone()

# def count_files(directory):
#     """计算指定目录下的文件总数"""
#     return sum(len(files) for _, _, files in os.walk(directory))

def main():
    parser = argparse.ArgumentParser(description='Generate image library based on hash tree.')
    parser.add_argument('--hash', required=True, help='Directory containing text files')
    parser.add_argument('--baseline', required=True, help='Root directory for baseline image library and SQLite database')
    parser.add_argument('--restore', required=True, help='Target directory to restore images')
    args = parser.parse_args()

    db_path = os.path.join(args.baseline, 'file_hashes.db')

    # 计算总文件数
    # total_files = count_files(args.hash)
    # with tqdm(total=total_files, desc='Processing files') as pbar:
    for root, dirs, files in os.walk(args.hash):
        for file in files:
            if file.endswith('.txt'):
                txt_file_path = os.path.join(root, file)
                photo_file_name = file[:-4] # Remove .txt extension

                # 检查photo_file_name_base是否以图片扩展名结尾
                if not any(photo_file_name.lower().endswith(ext) for ext in image_extensions):
                    # pbar.update(1)
                    continue  # 跳过非图片文件

                with open(txt_file_path, 'r') as f:
                    photo_hash = f.read().strip()

                result = get_photo_info(db_path, photo_hash)
                if not result:  # 如果没有找到对应的记录
                    print(f"No record found for hash {photo_hash} in {txt_file_path}")
                    # pbar.update(1)
                    continue

                zip_path, zip_name, file_path, file_name = result

                try:
                    zip_ref = zipfile.ZipFile(os.path.join(args.baseline, zip_path, zip_name), 'r')
                except zipfile.BadZipFile:
                    print(f"Error opening ZIP file: {zip_path}/{zip_name}")
                    # pbar.update(1)
                    continue


                # 假设 source_file 是通过 zip_path, zip_name, file_path, file_name 获取到的路径
                source_file = Path(file_path) / Path(file_name)
                source_file_str = str(source_file).replace(os.sep, '/')  # 确保路径使用正斜杠

                target_dir = os.path.join(args.restore, os.path.relpath(root, args.hash))
                os.makedirs(target_dir, exist_ok=True)
                target_file = os.path.join(target_dir, photo_file_name)

                try:
                    with zip_ref.open(source_file_str) as source, open(target_file, 'wb') as target:
                        target.write(source.read())
                except KeyError as e:
                    print(f"Error accessing file in ZIP: {e}")

                # pbar.update(1)

if __name__ == "__main__":
    main()